import torch
from torch import nn,optim
import C3D_model
from torch.utils.tensorboard import SummaryWriter
import os
from datetime import datetime
import socket
import timeit
from tqdm import tqdm
from  torch.utils.data import Dataset, DataLoader
from dataset import VideoDataset

def train_model(num_epochs,num_classes,lr,device,save_dir,train_dataloader,val_dataloader,test_dataloader):
    model =C3D_model.C3D(num_classes,pretrained=False)#C3D模型实例化

    #定义模型的损失函数，分类一般用交叉熵损失函数
    criterion=nn.CrossEntropyLoss()

    #定义优化器
    optimizer=optim.SGD(model.parameters(),lr=lr,momentum=0.9,weight_decay=5e-4)#随机梯度下降法

    #定义学习率的更新策略
    scheduler=optim.lr_scheduler.StepLR(optimizer,step_size=10,gamma=0.1)#假设lr=0.1,十轮之后0.01，再10轮之后0.001

    #将模型放到训练设备中(我搜的损失函数不用放到设备上）
    model.to(device)
    criterion.to(device)

    #日志记录
    log_dir=os.path.join(save_dir,'models',datetime.now().strftime("%b%d_%H-%M-%S")+'-'+socket.gethostname())
    writer = SummaryWriter(log_dir=log_dir)

    #开始模型的训练
    trainval_loaders={'train':train_dataloader,'val':val_dataloader}#验证集和训练集以字典的形式保存。这样可以每训练一轮，验证一轮
    trainval_sizes={x:len(trainval_loaders[x].dataset) for x in
                    ['train','val']}#计算训练和验证的大小{'train':8460,'val':2159}
    test_size=len(test_dataloader.dataset)#计算测试集的大小test_size:2701

    #开始训练
    for epoch in range(num_epochs):
        for phase in ['train','val']:
            start_time=timeit.default_timer()#计算训练开始时间
            running_loss=0.0#初始loss值
            running_corrects=0.0#初始正确值

            if phase=='train':
                model.train()
            else:
                model.eval()

            for inputs, labels in tqdm(trainval_loaders[phase]):#tqdm显示进度
                inputs=inputs.to(device)
                labels=labels.to(device)
                optimizer.zero_grad()#梯度清零

                if phase=='train':
                    outputs=model(inputs)
                else:
                    with torch.no_grad():
                        outputs=model(inputs)

                #计算softmax的输出概率
                probs=nn.Softmax(dim=1)(outputs)#在Conv2d的时候用argmax
                preds=torch.max(probs,dim=1)[1]#[1] 用于从元组中提取第二个张量，即最大概率值对应的类别索引
                labels=labels.long()#转化为long格式

                loss=criterion(outputs,labels)#计算损失函数

                if phase=='train':
                    loss.backward()
                    optimizer.step()#更新模型

                #计算该轮次所有loss值的累加
                running_loss+=loss.item()*inputs.size(0)

                #计算该轮次所有预测正确值的累加
                running_corrects+=torch.sum(preds==labels.data)
            scheduler.step()#更新lr
            epoch_loss=running_loss/trainval_sizes[phase]#计算该轮次的loss值，总loss值除以样本数量
            epoch_acc=running_corrects/trainval_sizes[phase]#就算该轮次的准确率值，总预测正确值除以样本数量

            if phase=='train':
                writer.add_scalar('data/train_loss_epoch',epoch_loss,epoch)#名字，精度，轮次
                writer.add_scalar('data/train_acc_epoch',epoch_acc,epoch)
            else:
                writer.add_scalar('data/val_loss_epoch', epoch_loss, epoch)  # 名字，精度，轮次
                writer.add_scalar('data/val_acc_epoch', epoch_acc, epoch)

            #计算停止的时间戳
            stop_time=timeit.default_timer()

            print("{} Epoch:{}/{} Loss:{} ACC:{}".format(phase,epoch+1,num_epochs,epoch_loss,epoch_acc))
            print("Execution Time:{}"+str(stop_time-start_time))
    writer.close()#关闭日志
    #保存训练好的权重(保存了所有轮次的模型，可以用以前Conv2d的代码，只保存最优的）
    torch.save({'epoch':epoch+1,'state_dict':model.state_dict(),'opt_dict':optimizer.state_dict(),},
               os.path.join(save_dir,'models','C3D'+'_epoch-'+str(epoch)+'.pth'))
    print("Save model at{}".format(os.path.join(save_dir,'models','C3D'+'_epoch-'+str(epoch)+'.pth.tar')))
    #tar是一种压缩文件

    #开始模型的测试
    model.eval()
    running_corrects=0.0#初始化准确的值
    for inputs,labels in tqdm(test_dataloader):#循环推理测试集中的数据，并计算准确率
        inputs=inputs.to(device)#将数据和标签放入到设备中
        labels=labels.long()
        labels=labels.to(device)
        with torch.no_grad():
            outputs=model(inputs)

        #计算softmax的输出概率
        probs=nn.Softmax(dim=1)(outputs)
        #计算最大概率值的标签
        preds=torch.max(probs,dim=1)[1]

        running_corrects+=torch.sum(preds==labels.data)
    epoch_acc=running_corrects.double()/test_size#计算准确率值
    print("Test Accuracy:{}".format(epoch_acc))

if __name__ == '__main__':
    #定义模型训练的设备
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    num_epochs=10
    num_classes=101
    lr=1e-3
    save_dir='model_result'#是一个文件夹

    train_data = VideoDataset(dataset_path='data/ucf101', images_path='train', clip_len=16)  # clip_len=16是取16张
    train_dataloader = DataLoader(train_data, batch_size=64, shuffle=True, num_workers=0)

    val_data = VideoDataset(dataset_path='data/ucf101', images_path='val', clip_len=16)  # clip_len=16是取16张
    val_dataloader = DataLoader(train_data, batch_size=64, shuffle=False, num_workers=0)

    train_data = VideoDataset(dataset_path='data/ucf101', images_path='test', clip_len=16)  # clip_len=16是取16张
    train_dataloader = DataLoader(train_data, batch_size=64, shuffle=False, num_workers=0)

    train_model(num_epochs, num_classes, lr, device, save_dir, train_dataloader, val_dataloader, test_dataloader)


